import matplotlib.pyplot as plt
import matplotlib.animation as animation


# 汉诺塔算法
def hanoi(n, source, target, auxiliary, moves):
    if n == 1:
        moves.append((source, target))
    else:
        hanoi(n - 1, source, auxiliary, target, moves)
        moves.append((source, target))
        hanoi(n - 1, auxiliary, target, source, moves)
    return moves


# 定义盘子数量
N = 3

# 初始化杆子和盘子
rods = {'A': list(range(N, 0, -1)), 'B': [], 'C': []}

# 获取汉诺塔移动步骤
moves = hanoi(N, 'A', 'C', 'B', [])

# 设置绘图
fig, ax = plt.subplots()
ax.set_xlim(0, 9)
ax.set_ylim(0, N + 2)
ax.set_xticks([2, 5, 8])
ax.set_xticklabels(['A', 'B', 'C'])
ax.set_yticks([])

# 绘制杆子
for pos in [2, 5, 8]:
    ax.plot([pos, pos], [0, N + 1], color='black', linewidth=5)

# 定义盘子的宽度和颜色
disk_width = [1 + 0.5 * i for i in range(N)]
disk_colors = plt.cm.viridis([0.2 * i for i in range(N)])

# 初始化显示的杆子和盘子
bars = []


def init():
    global bars
    bars = []
    for i in range(N):
        bar = ax.bar(2 - disk_width[i] / 2, 1, width=disk_width[i], color=disk_colors[i], align='edge')
        bars.append(bar)
    update_bars()
    return [b for bar in bars for b in bar]


# 更新显示的杆子和盘子
def update_bars():
    positions = {'A': 2, 'B': 5, 'C': 8}
    for rod in rods:
        for i, disk in enumerate(rods[rod]):
            bar = bars[disk - 1]
            bar[0].set_x(positions[rod] - disk_width[disk - 1] / 2)
            bar[0].set_y(i + 1)
            bar[0].set_height(0.8)


# 动画更新函数
def animate(i):
    source, target = moves[i]
    disk = rods[source].pop()
    rods[target].append(disk)
    update_bars()
    return [b for bar in bars for b in bar]


# 创建动画
ani = animation.FuncAnimation(fig, animate, frames=len(moves), init_func=init, blit=True, repeat=False)

plt.show()
